#!/usr/bin/env python3
import os
os.environ["ULTRALYTICS_LOGGING"] = "False"

from ultralytics.utils import LOGGER
import logging
from ultralytics import YOLO

try:
    LOGGER.remove()
except Exception:
    pass
logging.getLogger("ultralytics").setLevel(logging.ERROR)

import rospy
from sensor_msgs.msg import Image
from std_msgs.msg import Float32MultiArray
from cv_bridge import CvBridge
import cv2



class YOLODetector:
    def __init__(self):
        rospy.init_node('yolo_detector', anonymous=True)

        self.bridge = CvBridge()
        self.model = YOLO("/home/ycn/robot_arm_ws/src/OrbbecSDK_ROS1/models/best.pt", verbose=False)

        # 订阅彩色图像
        self.image_sub = rospy.Subscriber("/camera/color/image_raw", Image, self.callback)

        # 发布检测结果
        self.result_pub = rospy.Publisher("/yolo/detections", Float32MultiArray, queue_size=1)
        self.vis_pub = rospy.Publisher("/yolo/detections/image", Image, queue_size=1)

    def callback(self, msg):
        # 转为 OpenCV 图像
        cv_img = self.bridge.imgmsg_to_cv2(msg, "rgb8")

        # YOLO 推理，过滤置信度 < 0.7
        results = self.model(cv_img, conf=0.7)

        # 构建 Float32MultiArray
        data = Float32MultiArray()
        array = []

        # 绘制可视化图像
        vis_img = cv_img.copy()

        for r in results:
            boxes = r.boxes
            for i in range(len(boxes)):
                xyxy = boxes.xyxy[i].cpu().numpy()  # [xmin, ymin, xmax, ymax]
                conf = float(boxes.conf[i])
                cls = int(boxes.cls[i])

                cx = (xyxy[0] + xyxy[2]) / 2
                cy = (xyxy[1] + xyxy[3]) / 2
                w = xyxy[2] - xyxy[0]
                h = xyxy[3] - xyxy[1]

                array.extend([cls, conf, cx, cy, w, h])

                # 可视化：画矩形和类别/置信度
                cv2.rectangle(vis_img, (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3])), (0,255,0), 2)
                cv2.putText(vis_img, f"{cls}:{conf:.2f}", (int(xyxy[0]), int(xyxy[1]-5)),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 1)

        data.data = array
        self.result_pub.publish(data)

        # 发布可视化图像
        vis_msg = self.bridge.cv2_to_imgmsg(vis_img, "rgb8")
        self.vis_pub.publish(vis_msg)

    def run(self):
        rospy.spin()

if __name__ == "__main__":
    node = YOLODetector()
    node.run()
